'''
20250922 这个可用
train file :
    train_flux_kontext/train_flux_kontext_posenv_diffusers.sh
    train_flux_kontext/train_flux_kontext_posenv_diffusers.py
    
    --instance_data_dir /mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_processed/train/metadata.jsonl  \
    --image_column "file_name" \
    --cond_image_column "control_image" \
    --caption_column "prompt" \
    --aspect_ratio_buckets 720,1440 \
'''

import os
import json
from tqdm import tqdm
import shutil

# 原始路径配置
posenv_ori_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_ori'
posenv_tar_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_tar'
posenv_val_ori_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_val_ori'
posenv_val_tar_dir = '/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_val_tar'

# 新数据集路径
base_dir = "/mnt/nas/shengjie/datasets/KontextRefControl_poseenv_white_processed"
os.makedirs(base_dir, exist_ok=True)

# 创建子目录
train_image_dir = os.path.join(base_dir, "train/image")
train_control_dir = os.path.join(base_dir, "train/control_image")
val_image_dir = os.path.join(base_dir, "val/image")
val_control_dir = os.path.join(base_dir, "val/control_image")

os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_control_dir, exist_ok=True)
os.makedirs(val_image_dir, exist_ok=True)
os.makedirs(val_control_dir, exist_ok=True)

def process_dataset(ori_dir, tar_dir, output_meta_path, is_train=True):
    """处理单个数据集分区"""
    meta_data = []
    
    # 遍历目标图像
    for img_name in tqdm(os.listdir(tar_dir)):
        if not img_name.endswith('.png'):
            continue
            
        # 复制图像文件
        src_img = os.path.join(tar_dir, img_name)
        dst_img = os.path.join(train_image_dir if is_train else val_image_dir, img_name)
        shutil.copy2(src_img, dst_img)
        
        # 复制控制图像
        ctrl_img_name = img_name  # 假设同名
        src_ctrl = os.path.join(ori_dir, ctrl_img_name)
        dst_ctrl = os.path.join(train_control_dir if is_train else val_control_dir, ctrl_img_name)
        shutil.copy2(src_ctrl, dst_ctrl)
        
        # 读取提示词
        txt_path = os.path.join(tar_dir, img_name.replace('.png', '.txt'))
        prompt = open(txt_path, 'r').read().strip() if os.path.exists(txt_path) else ""
        
        # 添加到元数据
        meta_entry = {
            "file_name": f"image/{img_name}",
            "control_image": f"control_image/{img_name}",
            "prompt": prompt
        }
        meta_data.append(meta_entry)
    
    # 保存为JSONL
    with open(output_meta_path, 'w') as f:
        for item in meta_data:
            f.write(json.dumps(item) + '\n')

# 处理训练集
process_dataset(
    ori_dir=posenv_ori_dir,
    tar_dir=posenv_tar_dir,
    output_meta_path=os.path.join(base_dir, "train/metadata.jsonl"),
    is_train=True
)

# 处理验证集
process_dataset(
    ori_dir=posenv_val_ori_dir,
    tar_dir=posenv_val_tar_dir,
    output_meta_path=os.path.join(base_dir, "val/metadata.jsonl"),
    is_train=False
)

# 生成数据集描述文件
dataset_info = {
    "train_size": len(os.listdir(train_image_dir)),
    "val_size": len(os.listdir(val_image_dir)),
    "resolution": "720x1440",
    "default_prompt": "[refcontrolPoseEnv] change pose and env to photo with reference from left side,change person to photo with refence from right side."
}
with open(os.path.join(base_dir, "dataset_info.json"), 'w') as f:
    json.dump(dataset_info, f, indent=2)

print("✅ JSONL 数据集生成完成！")